
from typing import Any, List, MutableSequence, Optional, Tuple, Union

import torch

from pytorch_lightning.plugins.environments import TorchElasticEnvironment
from pytorch_lightning.utilities.exceptions import MisconfigurationException


def parse_npu_ids(npus: Optional[Union[int, str, List[int]]]) -> Optional[List[int]]:
    """
    Parses the NPU ids given in the format as accepted by the
    :class:`~pytorch_lightning.trainer.Trainer`.

    Args:
        npus: An int -1 or string '-1' indicate that all available NPUs should be used.
            A list of unique ints or a string containing list of comma separated unique integers
            indicates specific NPUs to use.
            An int 0 means that no NPUs should be used.
            Any int N > 0 indicates that NPUs [0..N) should be used.

    Returns:
        a list of npus to be used or ``None`` if no NPUs were requested

    Raises:
        MisconfigurationException:
            If no NPUs are available but the value of npus variable indicates request for NPUs
    """
    # Check that npus param is None, Int, String or List
    _check_data_type(npus)

    # Handle the case when no npus are requested
    if npus is None or (isinstance(npus, int) and npus == 0) or str(npus).strip() in ("0", "[]"):
        return None

    # We know user requested npus therefore if some of the
    # requested npus are not available an exception is thrown.
    npus = _normalize_parse_npu_string_input(npus)
    npus = _normalize_parse_npu_input_to_list(npus)
    if not npus:
        raise MisconfigurationException("npus requested but none are available.")
    if TorchElasticEnvironment.detect() and len(npus) != 1 and len(_get_all_available_npus()) == 1:
        # omit sanity check on torchelastic as by default shows one visible npus per process
        return npus

    # Check that npus are unique. Duplicate npus are not supported by the backend.
    _check_unique(npus)

    return _sanitize_npu_ids(npus)


def parse_cpu_cores(cpu_cores: Union[int, str, List[int]]) -> int:
    """Parses the cpu_cores given in the format as accepted by the ``devices`` argument in the
    :class:`~pytorch_lightning.trainer.Trainer`.

    Args:
        cpu_cores: An int > 0.

    Returns:
        an int representing the number of processes

    Raises:
        MisconfigurationException:
            If cpu_cores is not an int > 0
    """
    if isinstance(cpu_cores, str) and cpu_cores.strip().isdigit():
        cpu_cores = int(cpu_cores)

    if not isinstance(cpu_cores, int) or cpu_cores <= 0:
        raise MisconfigurationException("`devices` selected with `CPUAccelerator` should be an int > 0.")

    return cpu_cores


def _normalize_parse_npu_string_input(s: Union[int, str, List[int]]) -> Union[int, List[int]]:
    if not isinstance(s, str):
        return s
    if s == "-1":
        return -1
    if "," in s:
        return [int(x.strip()) for x in s.split(",") if len(x) > 0]
    return int(s.strip())


def _sanitize_npu_ids(npus: List[int]) -> List[int]:
    """Checks that each of the npus in the list is actually available. Raises a MisconfigurationException if any of
    the npus is not available.

    Args:
        npus: list of ints corresponding to npus indices

    Returns:
        unmodified npus variable

    Raises:
        MisconfigurationException:
            If machine has fewer available npus than requested.
    """
    all_available_npus = _get_all_available_npus()
    for npu in npus:
        if npu not in all_available_npus:
            raise MisconfigurationException(
                f"You requested npu: {npus}\n But your machine only has: {all_available_npus}"
            )
    return npus


def _normalize_parse_npu_input_to_list(npus: Union[int, List[int], Tuple[int, ...]]) -> Optional[List[int]]:
    assert npus is not None
    if isinstance(npus, (MutableSequence, tuple)):
        return list(npus)

    # must be an int
    if not npus:  # npus==0
        return None
    if npus == -1:
        return _get_all_available_npus()

    return list(range(npus))


def _get_all_available_npus() -> List[int]:
    """
    Returns:
         a list of all available npus
    """
    return list(range(torch.npu.device_count()))


def _check_unique(device_ids: List[int]) -> None:
    """Checks that the device_ids are unique.

    Args:
        device_ids: list of ints corresponding to npus indices

    Raises:
        MisconfigurationException:
            If ``device_ids`` of npus aren't unique
    """
    if len(device_ids) != len(set(device_ids)):
        raise MisconfigurationException("Device ID's (npus) must be unique.")


def _check_data_type(device_ids: Any) -> None:
    """Checks that the device_ids argument is one of: None, Int, String or List. Raises a MisconfigurationException
    otherwise.

    Args:
        device_ids: npu_cores parameter as passed to the Trainer

    Raises:
        MisconfigurationException:
            If ``device_ids`` of NPUs aren't ``int``, ``str``, sequence of ``int`` or ``None``
    """
    if device_ids is not None and (
        not isinstance(device_ids, (int, str, MutableSequence, tuple)) or isinstance(device_ids, bool)
    ):
        raise MisconfigurationException("Device ID's (NPU) must be int, string or sequence of ints or None.")

